import os
import numpy as np
import cv2
import re

def load_image(image_dir, row, col):
    pattern = re.compile(r'(.*)tr\d+-tc\d+\.png')
    for filename in os.listdir(image_dir):
        if re.search(f'tr{row}-tc{col}', filename):
            match = pattern.match(filename)
            if match:
                prefix = match.group(1)
                filepath = os.path.join(image_dir, filename)
                return cv2.imread(filepath), prefix, filename
    return None, None, None

def save_image(image, filename, output_dir):
    filepath = os.path.join(output_dir, filename)
    cv2.imwrite(filepath, image)
    print(f"Saved {filename} to {filepath}")

def combine_images(image_dir, image_output_dir):
    rows, cols = 37, 33
    image_tile_size = 1024
    flow_subgrid_rows, flow_subgrid_cols = 15, 13
    row_overlap, col_overlap = 4, 3

    os.makedirs(image_output_dir, exist_ok=True)

    prefix = None
    processed_subgrids = set()

    for subgrid_row in range(1, rows + 1, flow_subgrid_rows - row_overlap):
        if subgrid_row + flow_subgrid_rows - 1 > rows:
            subgrid_row = rows - flow_subgrid_rows + 1
        for subgrid_col in range(1, cols + 1, flow_subgrid_cols - col_overlap):
            if subgrid_col + flow_subgrid_cols - 1 > cols:
                subgrid_col = cols - flow_subgrid_cols + 1

            subgrid_id = (subgrid_row, subgrid_col)
            if subgrid_id in processed_subgrids:
                continue
            processed_subgrids.add(subgrid_id)

            image_combined = np.zeros((flow_subgrid_rows * image_tile_size, flow_subgrid_cols * image_tile_size, 3),
                                      dtype=np.uint8)

            for i in range(subgrid_row, subgrid_row + flow_subgrid_rows):
                for j in range(subgrid_col, subgrid_col + flow_subgrid_cols):
                    image_part, temp_prefix, filename = load_image(image_dir, i, j)
                    if prefix is None:
                        prefix = temp_prefix
                    row_offset = (i - subgrid_row) * image_tile_size
                    col_offset = (j - subgrid_col) * image_tile_size
                    image_combined[row_offset:row_offset + image_tile_size, col_offset:col_offset + image_tile_size,
                    :] = image_part

            save_image(image_combined,
                       f"{prefix}tr{(subgrid_row - 1) // (flow_subgrid_rows - row_overlap) + 1}-tc{(subgrid_col - 1) // (flow_subgrid_cols - col_overlap) + 1}.png",
                       image_output_dir)

# 示例用法
image_dir = 'E:\\wafer52\\11866_32nm'  # 替换为实际图片文件夹路径
image_output_dir = 'E:\\wafer52\\11866_32nm_big_combine'  # 替换为实际图片合并输出文件夹路径

combine_images(image_dir, image_output_dir)
